package wiki.xsx.filter;

import javax.servlet.*;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpServletResponseWrapper;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.OutputStreamWriter;
import java.io.PrintWriter;
import java.util.zip.GZIPOutputStream;

/**
 * Created by xsx on 2018/01/31.
 */
public class GZipFilter implements Filter {
    public void init(FilterConfig filterConfig) throws ServletException {

    }

    public void doFilter(ServletRequest servletRequest, ServletResponse servletResponse, FilterChain filterChain) throws IOException, ServletException {
        HttpServletResponse response = (HttpServletResponse) servletResponse;
        Wrapper wrapper = new Wrapper(response);
        filterChain.doFilter(servletRequest, wrapper);
        byte[] gzipData = gzip(wrapper.getResponseData());
        response.addHeader("Content-Encoding", "gzip");
        response.setContentLength(gzipData.length);
        ServletOutputStream output = servletResponse.getOutputStream();
        output.write(gzipData);
        output.flush();
    }

    public void destroy() {

    }

    private byte[] gzip(byte[] data) {
        ByteArrayOutputStream byteOutput = new ByteArrayOutputStream(10240);
        GZIPOutputStream output = null;

        try {
            output = new GZIPOutputStream(byteOutput);
            output.write(data);
        } catch (IOException e) {
            e.printStackTrace();
        } finally {
            if (output != null) {
                try {
                    output.close();
                } catch (IOException e) {
                }
            }
        }

        return byteOutput.toByteArray();
    }

    class Wrapper extends HttpServletResponseWrapper {
        public static final int OT_NONE = 0, OT_WRITER = 1, OT_STREAM = 2;
        private int outputType = OT_NONE;
        private ServletOutputStream output = null;
        private PrintWriter writer = null;
        private ByteArrayOutputStream buffer = null;

        public Wrapper(HttpServletResponse resp) throws IOException {
            super(resp);
            buffer = new ByteArrayOutputStream();
        }

        public PrintWriter getWriter() throws IOException {
            if (outputType == OT_STREAM)
                throw new IllegalStateException();
            else if (outputType == OT_WRITER)
                return writer;
            else {
                outputType = OT_WRITER;
                writer = new PrintWriter(new OutputStreamWriter(buffer, getCharacterEncoding()));
                return writer;
            }
        }

        public ServletOutputStream getOutputStream() throws IOException {
            if (outputType == OT_WRITER)
                throw new IllegalStateException();
            else if (outputType == OT_STREAM)
                return output;
            else {
                outputType = OT_STREAM;
                output = new WrappedOutputStream(buffer);
                return output;
            }
        }

        public void flushBuffer() throws IOException {
            if (outputType == OT_WRITER)
                writer.flush();

            if (outputType == OT_STREAM)
                output.flush();
        }

        public void reset() {
            outputType = OT_NONE;
            buffer.reset();
        }

        public byte[] getResponseData() throws IOException {
            flushBuffer();
            return buffer.toByteArray();
        }


        //内部类，扩展ServletOutputStream


        class WrappedOutputStream extends ServletOutputStream {

            private ByteArrayOutputStream buffer;

            public WrappedOutputStream(ByteArrayOutputStream buffer) {
                this.buffer = buffer;
            }

            public void write(int b) throws IOException {
                buffer.write(b);
            }

            public boolean isReady() {
                return false;
            }

            public void setWriteListener(WriteListener writeListener) {

            }
        }

    }
}
